from __future__ import print_function

import re
import time

import numpy as np
from bcc import BPF
from ctypes import *

import os
import sys
import fcntl
import dnslib
import argparse
import socket
import json
import struct
import ssl


def encode_dns(name):
    if len(name) + 1 > 255:
        raise Exception("DNS Name too long.")
    b = bytearray()
    for element in name.split('.'):
        sublen = len(element)
        if sublen > 63:
            raise ValueError('DNS label %s is too long' % element)
        b.append(sublen)
        b.extend(element.encode('ascii'))
    b.append(0)  # Add 0-len octet label for the root server
    return b


def add_cache_entry(cache, name):
    key = cache.Key()
    key_len = len(key.p)
    name_buffer = encode_dns(name)
    # Pad the buffer with null bytes if it is too short
    name_buffer.extend((0,) * (key_len - len(name_buffer)))
    key.p = (c_ubyte * key_len).from_buffer(name_buffer)
    leaf = cache.Leaf()
    leaf.p = (c_ubyte * 4).from_buffer(bytearray(4))
    cache[key] = leaf


parser = argparse.ArgumentParser(usage='For detailed information about usage,\
 try with -h option')
req_args = parser.add_argument_group("Required arguments")
req_args.add_argument("-i", "--interface", type=str, default="",
                      help="Interface name, defaults to all if unspecified.")
# req_args.add_argument("-d", "--domains", type=str, required=True, nargs="+",
#     help='List of domain names separated by space. For example: -d abc.def xyz.mno')
args = parser.parse_args()

# initialize BPF - load source code from http-parse-simple.c
bpf = BPF(src_file="dns.c", debug=0)
# print(bpf.dump_func("dns_test"))

function_dns_matching = bpf.load_func("dns_matching", BPF.SOCKET_FILTER)

BPF.attach_raw_socket(function_dns_matching, args.interface)

socket_fd = function_dns_matching.sock
fl = fcntl.fcntl(socket_fd, fcntl.F_GETFL)
fcntl.fcntl(socket_fd, fcntl.F_SETFL, fl & (~os.O_NONBLOCK))

so = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
host = "127.0.0.1"
port = 8887
CA_FILE = "/home/a/ou/key/ca.crt"
KEY_FILE = "/home/a/ou/key/client.key"
CERT_FILE = "/home/a/ou/key/client.crt"
context = ssl.SSLContext(ssl.PROTOCOL_TLS)
context.check_hostname = False
context.load_cert_chain(certfile=CERT_FILE, keyfile=KEY_FILE)
context.load_verify_locations(CA_FILE)
context.verify_mode = ssl.CERT_REQUIRED
sock = context.wrap_socket(so, server_side=False)
sock.connect((host, port))


class MyEncoder(json.JSONEncoder):
    def default(self, obj):
        if isinstance(obj, np.ndarray):
            return obj.tolist()
        elif isinstance(obj, bytes):
            return str(obj, encoding='utf-8')
        return json.JSONEncoder.default(self, obj)


temp = ''
while 1:
    # retrieve raw packet from socket
    try:
        packet_str = os.read(socket_fd, 2048)
    except KeyboardInterrupt:
        sock.close()
        sys.exit(0)
    packet_bytearray = bytearray(packet_str)

    ETH_HLEN = 14
    UDP_HLEN = 8
    # IP HEADER
    # calculate ip header length
    ip_header_length = packet_bytearray[ETH_HLEN]  # load Byte
    ip_header_length = ip_header_length & 0x0F  # mask bits 0..3
    ip_header_length = ip_header_length << 2  # shift to obtain length

    # calculate payload offset
    payload_offset = ETH_HLEN + ip_header_length + UDP_HLEN

    payload = packet_bytearray[payload_offset:]
    # pass the payload to dnslib for parsing
    dnsrec = dnslib.DNSRecord.parse(payload).questions
    # result = dnsrec.questions
    result = re.search(r'<DNS Question: (.*?) .*', str(dnsrec), re.M | re.I).group(1)[1:-2]
    if result != temp:
        domain = {'type': 'dns', 'domain name': result}
        json_dns = json.dumps(domain, cls=MyEncoder, indent=4).encode()
        data = struct.pack("I%ds" % (len(json_dns),), len(json_dns), json_dns)
        temp = result
        sock.send(data)
    # print (dnsrec, "\n")

